{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycocotools.coco\n",
    "import pycocotools.cocoeval\n",
    "import os\n",
    "import torch\n",
    "import PIL.Image\n",
    "import torchvision\n",
    "import torchvision.transforms\n",
    "import trt_pose.plugins\n",
    "import trt_pose.models\n",
    "import trt_pose.coco\n",
    "import torch2trt\n",
    "import tqdm\n",
    "import json\n",
    "from trt_pose.parse_objects import ParseObjects\n",
    "import torch2trt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tensorrt_module\n",
    "class My"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = trt_pose.models.dla34up_pose(18, 42).cuda().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('tasks/human_pose/experiments/dla34up_pose_256x256_A.json.checkpoints/epoch_249.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.zeros((1, 3, 256, 256)).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch2trt.tensorrt_converter('torch.split')\n",
    "def convert_split_dbg(ctx):\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "backbone_trt = torch2trt.torch2trt(model.backbone, [data], fp16_mode=True, max_workspace_size=1<<25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "backbone_torch = model.backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.backbone = backbone_trt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0012, device='cuda:0', grad_fn=<MaxBackward1>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.max(torch.abs(model.backbone(data) - backbone_trt(data)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> /home/john/.local/lib/python3.6/site-packages/torch2trt-0.0.0-py3.6.egg/torch2trt/torch2trt.py(171)mark_outputs()\n",
      "-> trt_tensor = torch_output._trt\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "(Pdb)  p torch_output.grad_fn\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<SplitWithSizesBackward object at 0x7fbff626a278>\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "(Pdb)  q\n"
     ]
    }
   ],
   "source": [
    "import pdb\n",
    "pdb.pm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.cuda().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap, paf = model(torch.zeros((1, 3, 256, 256)).cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 18, 128, 128])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cmap.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 42, 128, 128])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "paf.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_SHAPE = (256, 256)\n",
    "images_dir = 'val2017'\n",
    "annotation_file = 'annotations/person_keypoints_val2017_modified.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.21s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "cocoGtTmp = pycocotools.coco.COCO('annotations/person_keypoints_val2017_modified.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "topology = trt_pose.coco.coco_category_to_topology(cocoGtTmp.cats[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.24s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "cocoGt = pycocotools.coco.COCO('annotations/person_keypoints_val2017.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "catIds = cocoGt.getCatIds('person')\n",
    "imgIds = cocoGt.getImgIds(catIds=catIds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "parse_objects = ParseObjects(topology, cmap_threshold=0.05, link_threshold=0.1, cmap_window=11, line_integral_samples=7, max_num_parts=100, max_num_objects=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 / 2693\n",
      "100 / 2693\n",
      "200 / 2693\n",
      "300 / 2693\n",
      "400 / 2693\n",
      "500 / 2693\n",
      "600 / 2693\n",
      "700 / 2693\n",
      "800 / 2693\n",
      "900 / 2693\n",
      "1000 / 2693\n",
      "1100 / 2693\n",
      "1200 / 2693\n",
      "1300 / 2693\n",
      "1400 / 2693\n",
      "1500 / 2693\n",
      "1600 / 2693\n",
      "1700 / 2693\n",
      "1800 / 2693\n",
      "1900 / 2693\n",
      "2000 / 2693\n",
      "2100 / 2693\n",
      "2200 / 2693\n",
      "2300 / 2693\n",
      "2400 / 2693\n",
      "2500 / 2693\n",
      "2600 / 2693\n"
     ]
    }
   ],
   "source": [
    "results = []\n",
    "\n",
    "for n, imgId in enumerate(imgIds):\n",
    "    \n",
    "    # read image\n",
    "    img = cocoGt.imgs[imgId]\n",
    "    img_path = os.path.join(images_dir, img['file_name'])\n",
    "\n",
    "    image = PIL.Image.open(img_path).convert('RGB').resize(IMAGE_SHAPE)\n",
    "    data = transform(image).cuda()[None, ...]\n",
    "\n",
    "    cmap, paf = model(data)\n",
    "    cmap, paf = cmap.cpu(), paf.cpu()\n",
    "\n",
    "#     object_counts, objects, peaks, int_peaks = postprocess(cmap, paf, cmap_threshold=0.05, link_threshold=0.01, window=5)\n",
    "#     object_counts, objects, peaks = int(object_counts[0]), objects[0], peaks[0]\n",
    "    \n",
    "    object_counts, objects, peaks = parse_objects(cmap, paf)\n",
    "    object_counts, objects, peaks = int(object_counts[0]), objects[0], peaks[0]\n",
    "\n",
    "    for i in range(object_counts):\n",
    "        object = objects[i]\n",
    "        score = 0.0\n",
    "        kps = [0]*(17*3)\n",
    "        x_mean = 0\n",
    "        y_mean = 0\n",
    "        cnt = 0\n",
    "        for j in range(17):\n",
    "            k = object[j]\n",
    "            if k >= 0:\n",
    "                peak = peaks[j][k]\n",
    "                x = round(float(img['width'] * peak[1]))\n",
    "                y = round(float(img['height'] * peak[0]))\n",
    "                score += 1.0\n",
    "                kps[j * 3 + 0] = x\n",
    "                kps[j * 3 + 1] = y\n",
    "                kps[j * 3 + 2] = 2\n",
    "                x_mean += x\n",
    "                y_mean += y\n",
    "                cnt += 1\n",
    "\n",
    "        ann = {\n",
    "            'image_id': imgId,\n",
    "            'category_id': 1,\n",
    "            'keypoints': kps,\n",
    "            'score': score / 17.0\n",
    "        }\n",
    "        results.append(ann)\n",
    "    if n % 100 == 0:\n",
    "        print('%d / %d' % (n, len(imgIds)))\n",
    "#     break\n",
    "        \n",
    "with open('results.json', 'w') as f:\n",
    "    json.dump(results, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('results.json', 'w') as f:\n",
    "    json.dump(results, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading and preparing results...\n",
      "DONE (t=1.25s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "cocoDt = cocoGt.loadRes('results.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running per image evaluation...\n",
      "Evaluate annotation type *keypoints*\n",
      "DONE (t=7.26s).\n",
      "Accumulating evaluation results...\n",
      "DONE (t=0.12s).\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets= 20 ] = 0.263\n",
      " Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets= 20 ] = 0.547\n",
      " Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets= 20 ] = 0.219\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets= 20 ] = 0.218\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets= 20 ] = 0.339\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 20 ] = 0.331\n",
      " Average Recall     (AR) @[ IoU=0.50      | area=   all | maxDets= 20 ] = 0.586\n",
      " Average Recall     (AR) @[ IoU=0.75      | area=   all | maxDets= 20 ] = 0.314\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets= 20 ] = 0.230\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets= 20 ] = 0.471\n"
     ]
    }
   ],
   "source": [
    "cocoEval = pycocotools.cocoeval.COCOeval(cocoGt, cocoDt, 'keypoints')\n",
    "cocoEval.params.imgIds = imgIds\n",
    "cocoEval.params.catIds = [1]\n",
    "cocoEval.evaluate()\n",
    "cocoEval.accumulate()\n",
    "cocoEval.summarize()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
